from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session
from typing import List
from pydantic import BaseModel
from datetime import datetime

from ..database import get_db
from ..models.user import User
from ..models.knowledge_base import KnowledgeBase
from ..auth.auth import get_current_active_user

router = APIRouter()

# Pydantic 模式
class KnowledgeBaseCreate(BaseModel):
    name: str
    description: str = ""
    vector_db_type: str = "faiss"
    embedding_model: str = "text-embedding-ada-002"

class KnowledgeBaseUpdate(BaseModel):
    name: str = None
    description: str = None
    vector_db_type: str = None
    embedding_model: str = None

class KnowledgeBaseResponse(BaseModel):
    id: int
    name: str
    description: str
    vector_db_type: str
    embedding_model: str
    created_at: datetime
    updated_at: datetime
    owner_id: int
    
    class Config:
        from_attributes = True

@router.post("/", response_model=KnowledgeBaseResponse)
async def create_knowledge_base(
    kb: KnowledgeBaseCreate,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """创建知识库"""
    db_kb = KnowledgeBase(
        name=kb.name,
        description=kb.description,
        vector_db_type=kb.vector_db_type,
        embedding_model=kb.embedding_model,
        owner_id=current_user.id
    )
    db.add(db_kb)
    db.commit()
    db.refresh(db_kb)
    return db_kb

@router.get("/", response_model=List[KnowledgeBaseResponse])
async def list_knowledge_bases(
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """获取知识库列表"""
    knowledge_bases = db.query(KnowledgeBase).filter(
        KnowledgeBase.owner_id == current_user.id
    ).all()
    return knowledge_bases

@router.get("/{kb_id}", response_model=KnowledgeBaseResponse)
async def get_knowledge_base(
    kb_id: int,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """获取知识库详情"""
    kb = db.query(KnowledgeBase).filter(
        KnowledgeBase.id == kb_id,
        KnowledgeBase.owner_id == current_user.id
    ).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    return kb

@router.put("/{kb_id}", response_model=KnowledgeBaseResponse)
async def update_knowledge_base(
    kb_id: int,
    kb_update: KnowledgeBaseUpdate,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """更新知识库"""
    kb = db.query(KnowledgeBase).filter(
        KnowledgeBase.id == kb_id,
        KnowledgeBase.owner_id == current_user.id
    ).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    
    update_data = kb_update.dict(exclude_unset=True)
    for field, value in update_data.items():
        setattr(kb, field, value)
    
    db.commit()
    db.refresh(kb)
    return kb

@router.delete("/{kb_id}")
async def delete_knowledge_base(
    kb_id: int,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """删除知识库"""
    kb = db.query(KnowledgeBase).filter(
        KnowledgeBase.id == kb_id,
        KnowledgeBase.owner_id == current_user.id
    ).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    
    db.delete(kb)
    db.commit()
    return {"message": "知识库删除成功"} 